#!/usr/bin/env python


import numpy as np
from sklearn.cluster import KMeans
from sklearn.linear_model import LogisticRegression
from sklearn.svm import SVC
from sklearn.neighbors import KNeighborsClassifier
from sklearn.model_selection import train_test_split


from sklearn import datasets

digists = datasets.load_digits()
X_train, X_test, y_train, y_test = train_test_split(digists.data, digists.target, test_size=0.5)

def kmeans_pre(cls, n_prototypes=6):
    class cls_(cls):
        def fit(self, X, y, *args, **kwargs):
            kmeans = KMeans(n_clusters=n_prototypes)

            classes_ = np.unique(y_train)
            N, p = X.shape
            n_classes_ = classes_.size
            Ps = np.empty((n_prototypes*n_classes_, p))
            Cs = np.empty(n_prototypes*n_classes_)
            for c in classes_:
                X_c = X[y==c, :]
                kmeans.fit(X_c)  # do kmeans for each class
                Ps[c*n_prototypes:(c+1)*n_prototypes,:] = kmeans.cluster_centers_
                Cs[c*n_prototypes:(c+1)*n_prototypes] = np.full(n_prototypes, c)
            return super().fit(Ps, Cs, *args, **kwargs)
    return cls_


import pandas as pd

def compare(models:dict, X_train, X_test, y_train, y_test, timeit=False, predict=False):
    """Compare the performance of models
    models: machine learning models
    """
    
    from collections import defaultdict
    import time

    _result = defaultdict(list)

    for name, model in models.items():
        time1 = time.perf_counter()
        model.fit(X_train, y_train)
        _result['models'].append(name)
        if timeit:
            time2 = time.perf_counter()
            _result['times'].append(time2 - time1)
        _result['train scores'].append(model.score(X_train, y_train))
        _result['test scores'].append(model.score(X_test, y_test))
        if predict:
            _result[f'predict={y_test[0]}'].append(model.predict(X_test[:1])[0])
    result = pd.DataFrame(_result)
    print(result.round(decimals=4).to_markdown())

models = {
    'kmeans-model': kmeans_pre(KNeighborsClassifier)(),
    'model': KNeighborsClassifier()
}
compare(models, X_test, X_test, y_test, y_test, timeit=True)
